package com.hy.aicodegeneration.ai.guardrail;

import cn.hutool.core.util.StrUtil;
import com.hy.aicodegeneration.model.entity.PromptSafetyEvent;
import com.hy.aicodegeneration.ai.guardrail.service.PromptSafetyEventService ;
import com.hy.aicodeuser.service.UserService;
import com.hy.aicodegeneration.utils.SpringContextUtil;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.guardrail.InputGuardrail;
import dev.langchain4j.guardrail.InputGuardrailResult;
import jakarta.servlet.http.HttpServletRequest;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Component;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.time.LocalDateTime;
import java.util.*;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

@Component
@Slf4j
public class PromptSafetyInputGuardrail implements InputGuardrail {

    // 添加成员变量
    private final Long appId;

    // 添加构造函数
    public PromptSafetyInputGuardrail() {
        this(null); // 默认构造函数，用于向后兼容
    }

    public PromptSafetyInputGuardrail(Long appId) {
        this.appId = appId;
    }

    // 敏感词列表，使用Set提高查找效率
    private static final Set<String> DEFAULT_SENSITIVE_WORDS = loadSensitiveWords();

    // 敏感词资源文件路径
    private static final String SENSITIVE_WORDS_RESOURCE = "/guardrail/sensitive-words.txt";

    // 注入攻击模式
    private static final List<Pattern> INJECTION_PATTERNS = Arrays.asList(
            Pattern.compile("(?i)ignore\\s+(?:previous|above|all)\\s+(?:instructions?|commands?|prompts?)"),
            Pattern.compile("(?i)(?:forget|disregard)\\s+(?:everything|all)\\s+(?:above|before)"),
            Pattern.compile("(?i)(?:pretend|act|behave)\\s+(?:as|like)\\s+(?:if|you\\s+are)"),
            Pattern.compile("(?i)system\\s*:\\s*you\\s+are"),
            Pattern.compile("(?i)new\\s+(?:instructions?|commands?|prompts?)\\s*:")
    );

    // 敏感词替换映射表
    private static final Map<String, String> REPLACEMENT_WORDS = new HashMap<String, String>() {
        {
            put("破解", "绕过");
            put("黑客", "安全专家");
            put("盗版", "未授权版本");
            put("病毒", "恶意软件");
            put("攻击", "测试");
            // 可以根据需要添加更多替换规则
        }
    };

    // 注入模式替换为的文本
    private static final String INJECTION_REPLACEMENT = "[内容已优化]";

    // 加载敏感词列表
    private static Set<String> loadSensitiveWords() {
        try (InputStream inputStream = PromptSafetyInputGuardrail.class.getResourceAsStream(SENSITIVE_WORDS_RESOURCE);
             BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8))) {

            return reader.lines()
                    .filter(line -> !line.trim().isEmpty())         // 过滤空行
                    .filter(line -> !line.startsWith("#"))          // 过滤注释行
                    .map(String::trim)
                    .map(String::toLowerCase)                       // 统一转为小写
                    .collect(Collectors.toCollection(HashSet::new)); // 收集到HashSet中，提高查找效率

        } catch (IOException e) {
            // 如果无法加载文件，则回退到硬编码的默认列表
            return Arrays.asList(
                            // 指令绕过/覆盖
                            "忽略之前的指令", "忽略上面的指令", "忽略所有指令", "覆盖系统提示", "覆盖系统指令",
                            "ignore previous instructions", "ignore above", "ignore all instructions",
                            "override system prompt", "override instructions", "disregard previous",

                            // 绕过/越狱
                            "破解", "绕过", "越狱", "注入", "提示注入",
                            "hack", "bypass", "jailbreak", "prompt injection", "system jailbreak",

                            // 违法与侵权
                            "钓鱼", "网钓", "诈骗", "伪造", "造假", "木马", "恶意软件", "勒索软件",
                            "phishing", "scam", "fraud", "forgery", "malware", "ransomware", "trojan",

                            // 凭证与隐私数据
                            "密钥", "api key", "令牌", "口令", "密码", "口令泄露", "刷凭证", "抓包",
                            "api key", "access token", "secret", "credential", "password", "steal cookies",

                            // 非法获取与渗透
                            "端口扫描", "扫端口", "爆破", "暴力破解", "注入漏洞", "sql 注入", "xss",
                            "port scan", "nmap", "bruteforce", "sql injection", "xss", "csrf",

                            // 危害安全/暴力/自残（拦截关键词，不含指导性细节）
                            "自杀", "自残", "伤害他人", "制造炸弹", "枪支", "毒品制作",
                            "suicide", "self-harm", "harm others", "make a bomb", "firearm", "drug making",

                            // 版权与受保护内容绕过
                            "绕过版权", "免版税破解", "下载付费内容", "去水印", "解除 drm",
                            "bypass copyright", "pirated", "download paid content", "remove watermark", "bypass drm"
                    ).stream()
                    .map(String::toLowerCase)
                    .collect(Collectors.toCollection(HashSet::new));
        }
    }

    /**
     * 验证用户输入，并返回结果
     *
     * @param userMessage 用户输入的消息
     * @return 验证结果
     */
    @Override
    public InputGuardrailResult validate(UserMessage userMessage) {
        String input = userMessage.singleText();
        // 获取当前请求上下文（在主线程中）
        HttpServletRequest request = getCurrentRequest();

        // 初始化userId
        Long userId = null;
        if (request != null) {
            try {
                UserService userService = SpringContextUtil.getBean(UserService.class);
                userId = userService.getLoginUser(request).getId();
            } catch (Exception e) {
                // 用户未登录时忽略
                log.debug("用户未登录，无法获取用户ID");
            }
        }

        // 检查输入长度
        if (input.length() > 1000) {
            // 记录安全事件 - 使用新的辅助方法
            recordSafetyEvent(request, userId, this.appId, input, null, "LENGTH_LIMIT", null, "reject");
            // 抛出异常
            throw new RuntimeException("输入内容过长，不要超过 1000 字");
        }
        // 检查是否为空
        if (input.trim().isEmpty()) {
            // 记录安全事件 - 使用新的辅助方法
            recordSafetyEvent(request, userId, this.appId, input, null, "EMPTY_INPUT", null, "reject");
            throw new RuntimeException("输入内容不能为空");
        }

        // 执行内容净化
        PurificationResult result = purifyInput(input);

        // 如果需要重写
        if (result.needsRewriting()) {
            // 记录安全事件 - 使用新的辅助方法
            recordSafetyEvent(request, userId, this.appId, input, result.getRewrittenText(),
                    result.getRiskType(), result.getRiskKeywords(), "rewrite");

            // 由于langchain4j版本限制，我们创建重写后的消息但直接返回成功结果
            UserMessage rewrittenMessage = UserMessage.from(result.getRewrittenText());
            return InputGuardrailResult.success();
        }

        return InputGuardrailResult.success();
    }


    /**
     * 净化输入内容，替换敏感词和注入模式
     */
    private PurificationResult purifyInput(String input) {
        String rewrittenText = input;
        Set<String> riskKeywords = new HashSet<>();
        String riskType = null;
        boolean needsRewriting = false;

        // 检查敏感词并替换
        for (Map.Entry<String, String> entry : REPLACEMENT_WORDS.entrySet()) {
            if (rewrittenText.contains(entry.getKey())) {
                rewrittenText = rewrittenText.replace(entry.getKey(), entry.getValue());
                riskKeywords.add(entry.getKey());
                needsRewriting = true;
            }
        }

        // 检查注入攻击模式并替换
        for (Pattern pattern : INJECTION_PATTERNS) {
            if (pattern.matcher(rewrittenText).find()) {
                rewrittenText = pattern.matcher(rewrittenText).replaceAll(INJECTION_REPLACEMENT);
                riskType = "INJECTION_ATTEMPT";
                needsRewriting = true;
            }
        }

        // 检查是否包含需要完全拒绝的敏感词
        String lowerInput = input.toLowerCase();
        for (String sensitiveWord : DEFAULT_SENSITIVE_WORDS) {
            // 跳过已经在替换映射中的词，这些已经被处理过了
            if (REPLACEMENT_WORDS.containsKey(sensitiveWord)) {
                continue;
            }

            if (lowerInput.contains(sensitiveWord)) {
                // 对于完全拒绝的词，记录事件并返回fatal结果
                // 获取当前请求上下文（在主线程中）
                HttpServletRequest request = getCurrentRequest();

                // 获取userId
                Long userId = null;
                if (request != null) {
                    try {
                        UserService userService = SpringContextUtil.getBean(UserService.class);
                        userId = userService.getLoginUser(request).getId();
                    } catch (Exception e) {
                        // 用户未登录时忽略
                        log.debug("用户未登录，无法获取用户ID");
                    }
                }
                recordSafetyEvent(request, userId, this.appId, input, null, "CONTENT_RISK", sensitiveWord, "reject");
                return new PurificationResult(input, false, "CONTENT_RISK", sensitiveWord);
            }
        }

        // 如果有敏感词替换但没有注入模式，则设置风险类型
        if (needsRewriting && StrUtil.isEmpty(riskType)) {
            riskType = "CONTENT_MODIFIED";
        }

        return new PurificationResult(rewrittenText, needsRewriting, riskType,
                riskKeywords.isEmpty() ? null : String.join(",", riskKeywords));
    }

    /**
     * 新增辅助方法：在主线程中获取所有必要信息并构建完整事件对象
     */
    private void recordSafetyEvent(HttpServletRequest request, Long userId, Long appId,
                                   String originalPrompt, String rewrittenPrompt,
                                   String riskType, String riskKeywords, String processResult) {
        try {
            // 1. 在主线程中提前获取所有必要信息
            String ipAddress = getIpAddress(request);

            // 如果没有提供userId，尝试从当前请求获取
            if (userId == null && request != null) {
                try {
                    UserService userService = SpringContextUtil.getBean(UserService.class);
                    userId = userService.getLoginUser(request).getId();
                } catch (Exception e) {
                    // 用户未登录时忽略
                    log.debug("用户未登录，无法获取用户ID");
                }
            }

            // 2. 构建安全事件实体
            PromptSafetyEvent event = PromptSafetyEvent.builder()
                    .userId(userId != null ? userId : 0L) // 如果userId为null，使用默认值0
                    .appId(appId != null ? appId : 0L)    // 如果appId为null，使用默认值0
                    .originalPrompt(originalPrompt)
                    .rewrittenPrompt(rewrittenPrompt)
                    .riskType(riskType)
                    .riskKeywords(riskKeywords)
                    .processResult(processResult)
                    .ipAddress(ipAddress)
                    .createTime(LocalDateTime.now())
                    .updateTime(LocalDateTime.now())
                    .isDelete(0)
                    .build();

            // 3. 调用修改后的异步方法，直接传递完整的事件对象
            recordSafetyEventAsync(event);
        } catch (Exception e) {
            log.error("准备安全事件记录失败: {}", e.getMessage(), e);
        }
    }

    /**
     * 修改后的异步记录方法：直接接收完整的事件对象
     */
    @Async
    protected void recordSafetyEventAsync(PromptSafetyEvent event) {
        try {
            // 直接保存完整的事件对象
            PromptSafetyEventService eventService = SpringContextUtil.getBean(PromptSafetyEventService.class);
            boolean success = eventService.recordSafetyEvent(event);

            if (!success) {
                log.warn("记录安全事件到数据库失败");
            }
        } catch (Exception e) {
            log.error("记录安全事件失败: {}", e.getMessage(), e);
        }
    }

    /**
     * 获取当前请求
     */
    private HttpServletRequest getCurrentRequest() {
        try {
            return ((org.springframework.web.context.request.ServletRequestAttributes)
                    org.springframework.web.context.request.RequestContextHolder.getRequestAttributes())
                    .getRequest();
        } catch (Exception e) {
            return null;
        }
    }

    /**
     * 获取IP地址
     */
    private String getIpAddress(HttpServletRequest request) {
        if (request == null) {
            return "unknown";
        }

        String ip = request.getHeader("x-forwarded-for");
        if (StrUtil.isBlank(ip) || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("Proxy-Client-IP");
        }
        if (StrUtil.isBlank(ip) || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("WL-Proxy-Client-IP");
        }
        if (StrUtil.isBlank(ip) || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddr();
        }

        // 多级代理情况下，取第一个IP
        if (StrUtil.isNotBlank(ip) && ip.contains(",")) {
            ip = ip.split(",")[0].trim();
        }

        return ip;
    }

    /**
     * 净化结果内部类
     */
    private static class PurificationResult {
        private final String rewrittenText;
        private final boolean needsRewriting;
        private final String riskType;
        private final String riskKeywords;

        public PurificationResult(String rewrittenText, boolean needsRewriting,
                                  String riskType, String riskKeywords) {
            this.rewrittenText = rewrittenText;
            this.needsRewriting = needsRewriting;
            this.riskType = riskType;
            this.riskKeywords = riskKeywords;
        }

        public String getRewrittenText() {
            return rewrittenText;
        }

        public boolean needsRewriting() {
            return needsRewriting;
        }

        public String getRiskType() {
            return riskType;
        }

        public String getRiskKeywords() {
            return riskKeywords;
        }
    }
}